/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::cmp::Ordering;

use anyhow::Context;
use anyhow::Result;
use anyhow::bail;
use anyhow::ensure;
use anyhow::format_err;
use quickcheck::Arbitrary;
use quickcheck::Gen;
use rand_distr::Distribution;
use rand_distr::LogNormal;

use super::delta_apply::mpatch_fold;
use super::delta_apply::wrap_deltas;
use crate::errors::MononokeHgError;

#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Default)]
pub struct Delta {
    // Fragments should be in sorted order by start offset and should not overlap.
    frags: Vec<Fragment>,
}

impl Delta {
    /// Construct a new Delta object. Verify that `frags` is sane, sorted and
    /// non-overlapping.
    pub fn new(frags: Vec<Fragment>) -> Result<Self> {
        Self::verify(&frags)?;
        Ok(Delta { frags })
    }

    /// Construct a new Delta object given a fulltext (no delta).
    pub fn new_fulltext<T: Into<Vec<u8>>>(text: T) -> Self {
        Self {
            frags: vec![Fragment {
                start: 0,
                end: 0,
                content: text.into(),
            }],
        }
    }

    pub fn fragments(&self) -> &[Fragment] {
        self.frags.as_slice()
    }

    pub fn push(&mut self, frag: Fragment) {
        self.frags.push(frag);
    }

    /// If this delta might be a fulltext, return the fulltext. Note that we can only tell with
    /// certainty that something is *not* a fulltext. A delta with one fragment that inserts text
    /// in the beginning appears identical to a fulltext at this layer.
    pub fn maybe_fulltext(&self) -> Option<&[u8]> {
        if self.frags.len() == 1 && self.frags[0].start == 0 && self.frags[0].end == 0 {
            Some(self.frags[0].content.as_slice())
        } else {
            None
        }
    }

    fn verify(frags: &[Fragment]) -> Result<()> {
        let mut prev_frag: Option<&Fragment> = None;
        for (i, frag) in frags.iter().enumerate() {
            frag.verify().with_context(|| {
                MononokeHgError::InvalidFragmentList(format!("invalid fragment {}", i))
            })?;
            if let Some(prev) = prev_frag {
                if frag.start < prev.end {
                    let msg = format!(
                        "fragment {}: previous end {} overlaps with start {}",
                        i, prev.end, frag.start
                    );
                    bail!(MononokeHgError::InvalidFragmentList(msg));
                }
            }
            prev_frag = Some(frag);
        }
        Ok(())
    }
}

impl Arbitrary for Delta {
    fn arbitrary(g: &mut Gen) -> Self {
        let size = g.size();
        let nfrags = usize::arbitrary(g) % size;

        // Maintain invariants (start <= end, no overlap).
        let mut start = 0;
        let mut end = 0;

        let frags = (0..nfrags)
            .map(|_| {
                start = end + usize::arbitrary(g) % size;
                end = start + usize::arbitrary(g) % size;

                Fragment {
                    start,
                    end,
                    content: arbitrary_frag_content(g),
                }
            })
            .collect();
        Delta { frags }
    }

    fn shrink(&self) -> Box<dyn Iterator<Item = Self>> {
        // Not all instances generated by Vec::shrink will be
        // valid. Theoretically we could shrink in ways such that the invariants
        // are maintained, but just filtering is easier.
        Box::new(
            self.frags
                .shrink()
                .filter(|frags| Delta::verify(frags).is_ok())
                .map(|frags| Delta { frags }),
        )
    }
}

/// Represents a single contiguous modified region of text.
#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
pub struct Fragment {
    pub start: usize,
    pub end: usize,
    pub content: Vec<u8>,
}

impl Fragment {
    /// Return the end offset of this Fragment's content, after application.
    pub fn post_end(&self) -> usize {
        self.start + self.content.len()
    }

    // Return the length of the content
    pub fn content_length(&self) -> usize {
        self.content.len()
    }

    /// Return the change in text length this Fragment will cause when applied.
    pub fn length_change(&self) -> isize {
        self.content.len() as isize - (self.end - self.start) as isize
    }

    /// Return true if the given offset falls within this Fragment's content (post-application).
    pub fn contains_offset(&self, offset: usize) -> bool {
        self.start <= offset && offset < self.post_end()
    }

    fn verify(&self) -> Result<()> {
        if self.start > self.end {
            bail!("invalid fragment: start {} > end {}", self.start, self.end);
        }
        Ok(())
    }
}

impl Arbitrary for Fragment {
    fn arbitrary(g: &mut Gen) -> Self {
        let size = g.size();

        // Maintain invariant start <= end.
        let start = usize::arbitrary(g) % size;
        let end = start + usize::arbitrary(g) % size;

        Fragment {
            start,
            end,
            content: arbitrary_frag_content(g),
        }
    }

    fn shrink(&self) -> Box<dyn Iterator<Item = Self>> {
        Box::new(
            (self.start, self.end, self.content.clone())
                .shrink()
                .filter(|&(start, end, ref _content)| {
                    // shrink could produce bad values
                    start <= end
                })
                .map(|(start, end, content)| Fragment {
                    start,
                    end,
                    content,
                }),
        )
    }
}

struct GenRngWrapper<'a> {
    r#gen: &'a mut Gen,
}

impl<'a> rand::RngCore for GenRngWrapper<'a> {
    fn next_u32(&mut self) -> u32 {
        u32::arbitrary(self.r#gen)
    }

    fn next_u64(&mut self) -> u64 {
        u64::arbitrary(self.r#gen)
    }

    fn fill_bytes(&mut self, dest: &mut [u8]) {
        for b in dest {
            *b = u8::arbitrary(self.r#gen);
        }
    }

    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), rand::Error> {
        self.fill_bytes(dest);
        Ok(())
    }
}

fn arbitrary_frag_content(g: &mut Gen) -> Vec<u8> {
    let size = g.size();
    // Using a uniform distribution over size here can lead to extremely bloated
    // data structures. We also want to test zero-length data with more than a
    // (1/size) probability. So use a lognormal distribution.
    //
    // The choice of mean and stdev are pretty arbitrary, but they work well for
    // common sizes (~100).
    // TODO: make this more rigorous, e.g. by using params such that p95 = size.

    let lognormal = LogNormal::new(-3.0, 2.0).expect("Failed to make LogNormal");
    let content_len = ((size as f64) * lognormal.sample(&mut GenRngWrapper { r#gen: g })) as usize;

    let mut v = Vec::with_capacity(content_len);
    for b in v.iter_mut() {
        *b = u8::arbitrary(g);
    }
    v
}

/// Apply a Delta to an input text, returning the result.
pub fn apply(text: &[u8], delta: &Delta) -> Result<Vec<u8>> {
    let mut chunks = Vec::with_capacity(delta.frags.len() * 2);
    let mut off = 0;

    for frag in &delta.frags {
        ensure!(
            off <= frag.start,
            "Invalid delta, fragment start is less than current offset ({} < {})",
            frag.start,
            off
        );
        if off < frag.start {
            chunks.push(text.get(off..frag.start).ok_or_else(|| {
                format_err!(
                    "Invalid delta, the range {}..{} is out of bounds for provided text",
                    off,
                    frag.start
                )
            })?);
        }
        if !frag.content.is_empty() {
            chunks.push(frag.content.as_ref())
        }
        off = frag.end;
    }
    match off.cmp(&text.len()) {
        Ordering::Greater => bail!(
            "Invalid delta, fragment is referencing out of bounds content: {} > {}",
            off,
            text.len()
        ),
        Ordering::Less => chunks.push(&text[off..text.len()]),
        Ordering::Equal => {}
    };

    let size = chunks.iter().map(|c| c.len()).sum::<usize>();
    let mut output = Vec::with_capacity(size);
    for c in chunks {
        output.extend_from_slice(c);
    }
    Ok(output)
}

/// Apply a chain of Deltas to an input text, returning the result.
/// Pack all deltas into one delta, and apply a pack to input text.
pub fn apply_chain<I: IntoIterator<Item = Delta>>(text: &[u8], deltas: I) -> Result<Vec<u8>> {
    let mut res = Vec::from(text);

    let (wrapped_deltas, data) = wrap_deltas(deltas);

    if wrapped_deltas.is_empty() {
        Ok(res)
    } else {
        // fold all deltas into one delta using logarithmic algorithm
        let folded_wrapped_delta = mpatch_fold(&wrapped_deltas, 0, wrapped_deltas.len());

        // convert into Revlog Delta
        let folded_delta = folded_wrapped_delta.into_delta(data)?;

        // apply folded delta
        res = apply(&res, &folded_delta)?;
        Ok(res)
    }
}

/// XXX: Compatibility functions for the old bdiff module for testing purposes. The delta
/// module will replace that one once all instances of Vec\<bdiff::Delta\> are replaced
/// with delta::Delta, and this compatibility module will be removed at that time.
pub mod compat {
    use super::*;
    use crate::bdiff;

    pub fn convert<T>(deltas: T) -> Delta
    where
        T: IntoIterator<Item = bdiff::Delta>,
    {
        Delta {
            frags: deltas
                .into_iter()
                .map(|delta| Fragment {
                    start: delta.start,
                    end: delta.end,
                    content: delta.content,
                })
                .collect(),
        }
    }

    pub fn apply_deltas<T>(text: &[u8], deltas: T) -> Result<Vec<u8>>
    where
        T: IntoIterator<Item = Vec<bdiff::Delta>>,
    {
        apply_chain(text, deltas.into_iter().map(convert))
    }
}

#[cfg(test)]
mod tests {
    use mononoke_macros::mononoke;
    use quickcheck::quickcheck;

    use super::*;

    /// Test that fragments are verified properly.
    #[mononoke::test]
    fn test_delta_new() {
        #[rustfmt::skip]
        let test_cases = vec![
            (vec![Fragment { start: 0, end: 0, content: vec![] }], true),
            (vec![Fragment { start: 0, end: 5, content: vec![] }], true),
            (vec![Fragment { start: 0, end: 5, content: vec![] },
                  Fragment { start: 5, end: 8, content: vec![] }], true),
            (vec![Fragment { start: 0, end: 5, content: vec![] },
                  Fragment { start: 6, end: 9, content: vec![] }], true),
            (vec![Fragment { start: 0, end: 5, content: vec![] },
                  Fragment { start: 6, end: 5, content: vec![] }], false),
            (vec![Fragment { start: 0, end: 5, content: vec![] },
                  Fragment { start: 4, end: 8, content: vec![] }], false),
        ];

        for (frags, success) in test_cases.into_iter() {
            let delta = Delta::new(frags);
            if success {
                assert!(delta.is_ok());
            } else {
                assert!(delta.is_err());
            }
        }
    }

    #[mononoke::test]
    fn test_maybe_fulltext() {
        #[rustfmt::skip]
        let test_cases = vec![
            (vec![Fragment { start: 0, end: 0, content: vec![] }], true),
            (vec![Fragment { start: 0, end: 0, content: vec![b'a'] }], true),
            (vec![Fragment { start: 0, end: 1, content: vec![b'b'] }], false),
            (vec![Fragment { start: 1, end: 2, content: vec![b'c'] }], false),
            (vec![Fragment { start: 0, end: 0, content: vec![b'd'] },
                  Fragment { start: 1, end: 2, content: vec![b'e'] }], false),
        ];

        for (frags, maybe_fulltext) in test_cases.into_iter() {
            let delta = Delta::new(frags).unwrap();
            if maybe_fulltext {
                assert!(delta.maybe_fulltext().is_some());
            } else {
                assert!(delta.maybe_fulltext().is_none());
            }
        }
    }

    quickcheck! {
        fn delta_gen(delta: Delta) -> bool {
            Delta::verify(&delta.frags).is_ok()
        }

        fn delta_shrink(delta: Delta) -> bool {
            // This test is a bit redundant, but let's just verify.
            delta.shrink().take(100).all(|d| {
                Delta::verify(&d.frags).is_ok()
            })
        }

        fn fragment_gen(fragment: Fragment) -> bool {
            fragment.verify().is_ok()
        }

        fn fragment_shrink(fragment: Fragment) -> bool {
            fragment.shrink().take(100).all(|f| f.verify().is_ok())
        }
    }

    #[mononoke::test]
    fn test_apply_1() {
        let text = b"aaaa\nbbbb\ncccc\n";
        let delta = Delta {
            frags: vec![Fragment {
                start: 5,
                end: 10,
                content: (&b"xxxx\n"[..]).into(),
            }],
        };

        let res = apply(text, &delta).unwrap();
        assert_eq!(&res[..], b"aaaa\nxxxx\ncccc\n");
    }

    #[mononoke::test]
    fn test_apply_2() {
        let text = b"bbbb\ncccc\n";
        let delta = Delta {
            frags: vec![
                Fragment {
                    start: 0,
                    end: 5,
                    content: (&b"aaaabbbb\n"[..]).into(),
                },
                Fragment {
                    start: 10,
                    end: 10,
                    content: (&b"dddd\n"[..]).into(),
                },
            ],
        };

        let res = apply(text, &delta).unwrap();
        assert_eq!(&res[..], b"aaaabbbb\ncccc\ndddd\n");
    }

    #[mononoke::test]
    fn test_apply_3a() {
        let text = b"aaaa\nbbbb\ncccc\n";
        let delta = Delta {
            frags: vec![Fragment {
                start: 0,
                end: 15,
                content: (&b"zzzz\nyyyy\nxxxx\n"[..]).into(),
            }],
        };

        let res = apply(text, &delta).unwrap();
        assert_eq!(&res[..], b"zzzz\nyyyy\nxxxx\n");
    }

    #[mononoke::test]
    fn test_apply_3b() {
        let text = b"aaaa\nbbbb\ncccc\n";
        let delta = Delta {
            frags: vec![
                Fragment {
                    start: 0,
                    end: 5,
                    content: (&b"zzzz\n"[..]).into(),
                },
                Fragment {
                    start: 5,
                    end: 10,
                    content: (&b"yyyy\n"[..]).into(),
                },
                Fragment {
                    start: 10,
                    end: 15,
                    content: (&b"xxxx\n"[..]).into(),
                },
            ],
        };

        let res = apply(text, &delta).unwrap();
        assert_eq!(&res[..], b"zzzz\nyyyy\nxxxx\n");
    }

    #[mononoke::test]
    fn test_apply_4() {
        let text = b"aaaa\nbbbb";
        let delta = Delta {
            frags: vec![Fragment {
                start: 5,
                end: 9,
                content: (&b"bbbbcccc"[..]).into(),
            }],
        };

        let res = apply(text, &delta).unwrap();
        assert_eq!(&res[..], b"aaaa\nbbbbcccc");
    }

    #[mononoke::test]
    fn test_apply_5() {
        let text = b"aaaa\nbbbb\ncccc\n";
        let delta = Delta {
            frags: vec![Fragment {
                start: 5,
                end: 10,
                content: (&b""[..]).into(),
            }],
        };

        let res = apply(text, &delta).unwrap();
        assert_eq!(&res[..], b"aaaa\ncccc\n");
    }

    #[mononoke::test]
    fn test_malformed_1() {
        let text = b"aaaa";
        let delta = Delta {
            frags: vec![Fragment {
                start: 5,
                end: 10,
                content: (&b""[..]).into(),
            }],
        };

        match apply(text, &delta) {
            Err(_) => {}
            Ok(res) => panic!("Unexpected success: {:?}", res),
        }
    }

    #[mononoke::test]
    fn test_malformed_2() {
        let text = b"aaaa";
        let delta = Delta {
            frags: vec![Fragment {
                start: 0,
                end: 10,
                content: (&b""[..]).into(),
            }],
        };

        match apply(text, &delta) {
            Err(_) => {}
            Ok(res) => panic!("Unexpected success: {:?}", res),
        }
    }

    #[mononoke::test]
    fn test_apply_chain_logarithmic1() {
        let frags1 = vec![
            Fragment {
                start: 0,
                end: 1,
                content: (&b"xx"[..]).into(),
            },
            Fragment {
                start: 1,
                end: 2,
                content: (&b"xx"[..]).into(),
            },
        ];
        let delta1 = Delta::new(frags1).unwrap();

        let frags2 = vec![Fragment {
            start: 1,
            end: 3,
            content: (&b"bb"[..]).into(),
        }];
        let delta2 = Delta::new(frags2).unwrap();

        let deltas = vec![delta1, delta2];
        let text = b"aaaaaa";

        let res = apply_chain(text, deltas).unwrap();
        assert_eq!(&res[..], b"xbbxaaaa");
    }

    #[mononoke::test]
    fn test_apply_chain_logarithmic2() {
        let frags1 = vec![Fragment {
            start: 1,
            end: 2,
            content: (&b"xx"[..]).into(),
        }];
        let delta1 = Delta::new(frags1).unwrap();

        let deltas = vec![delta1];
        let text = b"aaaaaa";

        let res = apply_chain(text, deltas).unwrap();
        assert_eq!(&res[..], b"axxaaaa");
    }

    #[mononoke::test]
    fn test_apply_chain_logarithmic3() {
        let frags1 = vec![
            Fragment {
                start: 5,
                end: 10,
                content: (&b"xxx"[..]).into(),
            },
            Fragment {
                start: 10,
                end: 12,
                content: (&b"yyyyy"[..]).into(),
            },
        ];
        let delta1 = Delta::new(frags1).unwrap();

        let frags2 = vec![Fragment {
            start: 6,
            end: 7,
            content: (&b"zzzzzzz"[..]).into(),
        }];
        let delta2 = Delta::new(frags2).unwrap();

        let deltas = vec![delta1, delta2];
        let text = b"aaaaabbbbbccccc";

        let res = apply_chain(text, deltas).unwrap();
        assert_eq!(&res[..], b"aaaaaxzzzzzzzxyyyyyccc");
    }

    #[mononoke::test]
    fn test_apply_chain_logarithmic_append() {
        let frags1 = vec![Fragment {
            start: 1,
            end: 1,
            content: (&b"xxx"[..]).into(),
        }];
        let delta1 = Delta::new(frags1).unwrap();

        let frags2 = vec![Fragment {
            start: 4,
            end: 4,
            content: (&b"zzz"[..]).into(),
        }];
        let delta2 = Delta::new(frags2).unwrap();

        let deltas = vec![delta1, delta2];
        let text = b"a";

        let res = apply_chain(text, deltas).unwrap();

        assert_eq!(&res[..], b"axxxzzz");
    }
}
